import lib.tools as my_tools
import sys
import matplotlib
# 使用非交互式后端，避免GUI线程问题
matplotlib.use('Agg')  # 使用Agg后端生成图片文件，不显示GUI
import matplotlib.pyplot as plt
import numpy as np
import os
import concurrent.futures
import tqdm

def process_batch(data_set, batch_idx, batch_size, num_items, num_batches, pbar):
    start_idx = batch_idx * batch_size
    end_idx = min((batch_idx + 1) * batch_size, num_items)
    
    # 获取当前批次数据
    batch_data = data_set[start_idx:end_idx]
    batch_indices = range(start_idx, end_idx)
    
    D_value = batch_data[:, 0]
    omega_x = batch_data[:, 1]
    omega_y = batch_data[:, 2]
    desired_D_value = batch_data[:, 10]

    # 为每个批次创建单独的子文件夹
    batch_dir = f'check_result/batch_{batch_idx:04d}'
    os.makedirs(batch_dir, exist_ok=True)

    # 使用Agg后端，无需担心多线程冲突
    # 绘制D-value和desired D-value的散点图（仅显示前120个点用于预览）
    fig = plt.figure(figsize=(12, 8))
    plt.scatter(batch_indices[:120], D_value[:120], s=5, label='D_value')
    plt.scatter(batch_indices[:120], desired_D_value[:120], alpha=1, s=2, label='desired D_value')
    plt.xlabel('Index')
    plt.ylabel('Value')
    plt.legend(loc='upper right')
    plt.xlim([start_idx, start_idx + 120])
    plt.grid(True)
    plt.title(f'D_value and desired D_value (Batch {batch_idx})')
    fig.savefig(f'{batch_dir}/attitude_D_value_scatter.png')
    plt.close(fig)

    # 绘制完整的D_value曲线
    fig = plt.figure(figsize=(20, 10))
    plt.scatter(batch_indices, D_value, s=0.5)
    plt.grid(True)
    plt.title(f'D_value Curve (Batch {batch_idx})')
    fig.savefig(f'{batch_dir}/D_value.png')
    plt.close(fig)

    # 绘制Omega_x曲线
    fig = plt.figure(figsize=(20, 10))
    plt.scatter(batch_indices, omega_x, s=0.5)
    plt.grid(True)
    plt.title(f'Omega_x Curve (Batch {batch_idx})')
    fig.savefig(f'{batch_dir}/Omega_x.png')
    plt.close(fig)

    # 绘制Omega_y曲线
    fig = plt.figure(figsize=(20, 10))
    plt.scatter(batch_indices, omega_y, s=0.5)
    plt.grid(True)
    plt.title(f'Omega_y Curve (Batch {batch_idx})')
    fig.savefig(f'{batch_dir}/Omega_y.png')
    plt.close(fig)

    pbar.set_description(f"Processing batch {batch_idx + 1}/{num_batches} ({end_idx - start_idx} items)")
    pbar.update(1)

def check_attitude():
    ## npy文件格式：
    datasetfilename = 'sorted_dataset_attitude.npy'
    file_size = os.path.getsize(datasetfilename)
    num_items = file_size // (12 * 4)  # 12个float32，每个4字节
    data_set = np.memmap(datasetfilename, dtype='float32', mode='r', shape=(num_items,12))

    # 创建check_result目录
    os.makedirs('check_result', exist_ok=True)
    
    # 批次大小和线程数配置
    batch_size = 100000
    num_batches = (num_items + batch_size - 1) // batch_size
    max_workers = min(8, os.cpu_count())  # 最多使用8个线程

    # 使用线程池处理批次
    with tqdm.tqdm(total=num_batches, desc="Processing batches") as pbar:
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = []
            for batch_idx in range(num_batches):
                # 传递num_batches参数和进度条
                futures.append(executor.submit(process_batch, data_set, batch_idx, batch_size, num_items, num_batches, pbar))
            
            # 等待所有任务完成
            concurrent.futures.wait(futures)
            for future in futures:
                if future.exception():
                    print(f"Error processing batch: {future.exception()}")



COMMANDS = {
    "attitude": "Check attitude dataset"
}

if __name__ == "__main__":

    if len(sys.argv) == 1:
        my_tools.print_help(sys.argv[0], COMMANDS)

    if len(sys.argv) == 2:
        
        # 针对不同选项启用不同的训练
        if sys.argv[1] == "attitude":
            check_attitude()